
import { Context } from 'koa';
import { GenID } from '../core/utils';
import * as fs from 'fs';
import * as path from 'path';
import { GenPosterGraph } from '../workflows/genPosterWork/genPoster';
import { DBService } from '../core/services/dbService';
import { IWorkflow } from '../core';
import { ImageEditorGraph } from '../workflows/imageEditorWork/imageEditor';


const sendRes = (ctx: Context, result: any, code?: number) => {
    result || (result = {});
    ctx.body = {
        code: code || 0,
        payload: result,
    };
}
const msgSvr = DBService.StoreMessageService;
const ensureUploadDir = () => {
    let uploadDir = path.join(__dirname, '../../uploads');
    if (process.env.USER_DIR) {
        uploadDir = path.join(process.env.USER_DIR, 'uploads');
    }

    if (!fs.existsSync(uploadDir)) {
        fs.mkdirSync(uploadDir, { recursive: true });
    }
    return uploadDir;
};

const getImageContentType = (filename: string): string => {
    const ext = path.extname(filename).toLowerCase();
    const mimeTypes: { [key: string]: string } = {
        '.jpg': 'image/jpeg',
        '.jpeg': 'image/jpeg',
        '.png': 'image/png',
        '.gif': 'image/gif',
        '.bmp': 'image/bmp',
        '.webp': 'image/webp',
        '.svg': 'image/svg+xml'
    };
    return mimeTypes[ext] || 'application/octet-stream';
};


export async function WorkflowRun(ctx: Context) {
    const userId = ctx.headers["userid"];
    let convID = ctx.headers["convid"];
    const reqParams = ctx.request.body;

    const { workflowId } = ctx.params;
    let { isStream } = ctx.request.query as any;
    if (isStream == undefined) {
        isStream = true;
    }

    if (isStream == "false") {
        isStream = false;
    }
    if (isStream == "true") {
        isStream = true;
    }

    if (!userId) {
        throw new Error("userId is required");
    }

    if (!convID) {
        convID = GenID();
    }

    let newSvr: IWorkflow;
    if (workflowId == "GenPosterGraph") {
        newSvr = new GenPosterGraph();
    }
    else if (workflowId == "ImageEditorGraph") {
        newSvr = new ImageEditorGraph();
    }

    if (!newSvr) {
        throw new Error("Invalid workflowId");
    }

    if (isStream) {
        ctx.res.setHeader('Content-Type', 'text/event-stream');
        ctx.res.setHeader('Cache-Control', 'no-cache');
        ctx.res.setHeader('Connection', 'keep-alive');
        ctx.res.statusCode = 200;

        newSvr.OnLLmMessage = (params) => {
            ctx.res.write(params);
        };

        newSvr.OnLLmEnd = (args) => {
            ctx.res.write(`data:${JSON.stringify(args)}\n\n`);
            ctx.res.write("data:[DONE]");
            ctx.res.end();
        }
    }

    const inputParams = {
        convID,
        userID: userId,
        ...reqParams
    }
    const result = await newSvr.Run(inputParams);
    if (!isStream) {
        ctx.response.body = {
            content: result
        };
    }
}

export async function WorkflowStop(ctx: Context) {
    const { workflowId } = ctx.params;
    let newSvr: IWorkflow = new GenPosterGraph();
    newSvr.Stop();

    ctx.response.body = { success: true };
}

export async function QueryConvs(ctx: Context) {
    const result = await msgSvr.QueryConvs({});
    ctx.response.body = result;
}

export async function GetMessages(ctx: Context) {
    let { convId } = ctx.params;
    ctx.response.body = await msgSvr.GetMessagesByConvId(convId);
}

export async function CleanMessages(ctx: Context) {
    let { convId } = ctx.params;
    ctx.response.body = await msgSvr.CleanMessageByConvId(convId);
}

/**
 * 图片上传接口
 * @param ctx Koa Context
 */
export async function UploadImage(ctx: Context) {
    try {
        // 确保上传目录存在
        const uploadDir = ensureUploadDir();

        // 检查是否有文件上传
        if (!ctx.request.files || !ctx.request.files.image) {
            ctx.status = 400;
            ctx.body = {
                code: 400,
                message: 'No file uploaded'
            };
            return;
        }

        const file: any = Array.isArray(ctx.request.files.image)
            ? ctx.request.files.image[0]
            : ctx.request.files.image;

        // 检查文件类型是否为图片
        if (!file.mimetype.startsWith('image/')) {
            ctx.status = 400;
            ctx.body = {
                code: 400,
                message: 'Only image files are allowed'
            };
            // 删除已保存的文件
            if (fs.existsSync(file.path)) {
                fs.unlinkSync(file.path);
            }
            return;
        }

        // 生成唯一文件名
        const newPath = path.join(uploadDir, file.newFilename);

        // 移动文件到指定目录
        fs.copyFileSync(file.filepath, newPath);
        fs.unlinkSync(file.filepath);

        // 构造返回结果
        const result = {
            filename: file.newFilename,
            originalname: file.name,
            mimetype: file.type,
            size: file.size,
            // url: `/uploads/${file.newFilename}`,
            // path: newPath
        };

        sendRes(ctx, result);
    } catch (error: any) {
        console.error('Upload image error:', error);
        ctx.status = 500;
        ctx.body = {
            code: 500,
            message: error.message || 'Upload failed'
        };
    }
}

export async function GetImage(ctx: Context) {
    const { filename } = ctx.params;
    const uploadDir = ensureUploadDir();
    const filePath = path.join(uploadDir, filename);
    if (!fs.existsSync(filePath)) {
        ctx.status = 404;
        ctx.body = {
            code: 404,
            message: 'File not found'
        };
        return;
    }
    ctx.set('Content-Type', getImageContentType(filename));
    ctx.body = fs.createReadStream(filePath);
}
